package com.kingsoft.dc.khaos.extender.auth;

import com.alibaba.fastjson.JSON;
import com.kingsoft.dc.khaos.extender.auth.utils.ConfigUtils;
import com.kingsoft.dc.khaos.extender.auth.utils.SqlParser;
import com.kingsoft.dc.khaos.extender.meta.model.AuthTable;
import com.kingsoft.dc.khaos.extender.model.DDLType;
import com.kingsoft.dc.khaos.extender.model.ParamConfig;
import org.apache.commons.io.FileUtils;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.nio.charset.Charset;
import java.util.*;
import java.util.stream.Collectors;

/**
 * Created by jing on 19/8/6.
 */
public class SqlAuthentication extends DmAuthentication {
    private Logger logger = LoggerFactory.getLogger(this.getClass());

    public boolean authenticate(String sql, String dbType) {
        ParamConfig paramConfig = ConfigUtils.loadParamConfigFromEnv();
        List<AuthTable> tableList = extractAuthTables(sql, dbType);
        return checkPriv(tableList, paramConfig);
    }

    public boolean authenticate(String sql, String dbType, String config) {
        ParamConfig paramConfig = JSON.parseObject(config, ParamConfig.class);
        List<AuthTable> tableList = extractAuthTables(sql, dbType);
        return checkPriv(tableList, paramConfig);
    }

    public boolean authenticate(String sql, ParamConfig paramConfig) {
//        ParamConfig paramConfig = JSON.parseObject(config, ParamConfig.class);
        List<AuthTable> tableList = extractAuthTables(sql, paramConfig.getDbType());
        return checkPriv(tableList, paramConfig);
    }

    /**
     * sparksql 新接口,待过滤功能的鉴权,authTables传要过滤的库表及action
     * @param sql
     * @param paramConfig
     * @param authTables
     * @return
     */
    public boolean authenticate(String sql, ParamConfig paramConfig, List<AuthTable> authTables) {
//        ParamConfig paramConfig = JSON.parseObject(config, ParamConfig.class);

        List<AuthTable> tableList = extractAuthTables(sql, paramConfig.getDbType());
        logger.info("paramConfig.getDbType()_{}",paramConfig.getDbType());
        logger.info("paramConfig.getDbType()_sql:{}",sql);

        for (AuthTable filter : authTables) {
            logger.info(String.format("Filter AuthTable: db:%s, table:%s, actions:[%s]", filter.getDbName(), filter.getTblName(), Arrays.toString(filter.getPrivSet())));
//            for (AuthTable table : tableList.iterator()) {
            for (int i = 0; i < tableList.size(); i++) {
                AuthTable table = tableList.get(i);
                logger.info(String.format("AuthTable: db:%s, table:%s, actions:[%s]", table.getDbName(), table.getTblName(), Arrays.toString(table.getPrivSet())));
                if (StringUtils.isNotBlank(filter.getDbName())
                        && StringUtils.isNotBlank(filter.getTblName())
                        && filter.getDbName().equalsIgnoreCase(table.getDbName())
                        && filter.getTblName().equalsIgnoreCase(table.getTblName())) {

                    List<String> actionList = Arrays.stream(table.getPrivSet()).collect(Collectors.toList());
                    List<String> filterActionList = Arrays.stream(filter.getPrivSet()).collect(Collectors.toList());
                    actionList.removeAll(filterActionList);
                    if (actionList.size() == 0) {
                        tableList.remove(i);
                        break;
                    }
                    table.setPrivSet(actionList.toArray(new String[0]));
                    logger.info(String.format("Check Table: db:%s, table:%s, actions:", table.getDbName(), table.getTblName(), Arrays.toString(table.getPrivSet())));
                    break;
                }
            }
        }
        if (tableList.size() == 0) return true;
        return checkPriv(tableList, paramConfig);
    }

    /**
     * sparksql 新接口,待过滤功能的鉴权,authTables传要过滤的库表及action
     * @param sql
     * @param paramConfig
     * @param authTables
     * @return
     */
    public boolean authenticate(String sql, ParamConfig paramConfig, List<AuthTable> authTables,String dbName) {
//        ParamConfig paramConfig = JSON.parseObject(config, ParamConfig.class);
        logger.info("paramConfig.getDbType():",dbName);
        List<AuthTable> tableList = extractAuthTables(sql, paramConfig.getDbType(),dbName,authTables);
        logger.info("paramConfig.getDbType():",paramConfig.getDbType());
        logger.info("paramConfig.getDbType()_sql:{}",sql);
        logger.info("paramConfig.getDbType()_dbName:",dbName);
        for (AuthTable filter : authTables) {
            logger.info(String.format("Filter AuthTable: db:%s, table:%s, actions:[%s]", filter.getDbName(), filter.getTblName(), Arrays.toString(filter.getPrivSet())));
//            for (AuthTable table : tableList.iterator()) {
            for (int i = 0; i < tableList.size(); i++) {
                AuthTable table = tableList.get(i);
                logger.info(String.format("AuthTable: db:%s, table:%s, actions:[%s]", table.getDbName(), table.getTblName(), Arrays.toString(table.getPrivSet())));
                if (StringUtils.isNotBlank(filter.getDbName())
                        && StringUtils.isNotBlank(filter.getTblName())
                        && filter.getDbName().equalsIgnoreCase(table.getDbName())
                        && filter.getTblName().equalsIgnoreCase(table.getTblName())) {

                    List<String> actionList = Arrays.stream(table.getPrivSet()).collect(Collectors.toList());
                    List<String> filterActionList = Arrays.stream(filter.getPrivSet()).collect(Collectors.toList());
                    actionList.removeAll(filterActionList);
                    if (actionList.size() == 0) {
                        tableList.remove(i);
                        break;
                    }
                    table.setPrivSet(actionList.toArray(new String[0]));
                    logger.info(String.format("Check Table: db:%s, table:%s, actions:", table.getDbName(), table.getTblName(), Arrays.toString(table.getPrivSet())));
                    break;
                }
            }
        }
        if (tableList.size() == 0) return true;
        return checkPriv(tableList, paramConfig);
    }

    /**
     * 获取要检测的table列表，table信息包含数据库名，数据表名，和操作类型
     *
     * @return
     */
    public List<AuthTable> extractAuthTables(String sql, String dbType, boolean enableDDL) {
        List<AuthTable> tableAuths = new ArrayList<>();
        Map<String, TreeSet<String>> authTables = SqlParser.getAuthTables(sql, dbType);
        for (Map.Entry<String, TreeSet<String>> entry : authTables.entrySet()) {
            String fullTableName = entry.getKey();
            logger.info("fullTableName-fullTableName:{}",fullTableName);//DEFAULT.lyhivesink
            TreeSet<String> ddlType = entry.getValue();
            logger.info("ddlType:{}",ddlType);//[SELECT]
            List<String> ddlTypes = DDLType.getNames();
            logger.info("ddlTypes:{}",ddlTypes);//[CREATE, DROP, ALTER, CREATE_INDEX, DROP_INDEX]
            ddlTypes.retainAll(ddlType);
            logger.info("ddlTypes.retainAll(ddlType):{}",ddlTypes.retainAll(ddlType));//false
            if (!ddlTypes.isEmpty() && !enableDDL) {
                logger.error("不能包含DDL操作:" + StringUtils.join(ddlTypes.toArray(), ","));
                tableAuths.clear();
                break;
            } else {
                AuthTable authTable = new AuthTable(fullTableName, ddlType.toArray(new String[ddlType.size()]));
                tableAuths.add(authTable);
            }
        }
        return tableAuths;
    }

    /**
     * mppsql和sparksql获取要检测的table列表，table信息包含数据库名，数据表名，和操作类型
     *
     * @return
     */
    public List<AuthTable> extractAuthTables(String sql, String dbType, boolean enableDDL,String dbName,List<AuthTable> viewAuthTables) {
        List<AuthTable> tableAuths = new ArrayList<>();
        Map<String, TreeSet<String>> authTables = SqlParser.getAuthTables(sql, dbType, dbName,viewAuthTables);
        for (Map.Entry<String, TreeSet<String>> entry : authTables.entrySet()) {
            String fullTableName = entry.getKey();
            logger.info("fullTableName-fullTableName:{}",fullTableName);//DEFAULT.lyhivesink
            TreeSet<String> ddlType = entry.getValue();
            logger.info("ddlType:{}",ddlType);//[SELECT]
            List<String> ddlTypes = DDLType.getNames();
            logger.info("ddlTypes:{}",ddlTypes);//[CREATE, DROP, ALTER, CREATE_INDEX, DROP_INDEX]
            ddlTypes.retainAll(ddlType);
            logger.info("ddlTypes.retainAll(ddlType):{}",ddlTypes.retainAll(ddlType));//false
            if (!ddlTypes.isEmpty() && !enableDDL) {
                logger.error("不能包含DDL操作:" + StringUtils.join(ddlTypes.toArray(), ","));
                tableAuths.clear();
                break;
            } else {
                AuthTable authTable = new AuthTable(fullTableName, ddlType.toArray(new String[ddlType.size()]));
                tableAuths.add(authTable);
            }
        }
        return tableAuths;
    }
    /**
     * 获取要检测的table列表，table信息包含数据库名，数据表名，和操作类型
     *
     * @return
     */
    public List<AuthTable> extractAuthTables(String sql, String dbType) {
        logger.info("extractAuthTables-sql:{}",sql);
        return extractAuthTables(sql, dbType, false);
    }

    /**
     * 获取要检测的table列表，table信息包含数据库名，数据表名，和操作类型
     *
     * @return
     */
    public List<AuthTable> extractAuthTables(String sql, String dbType,String dbName,List<AuthTable> viewAuthTables) {
        logger.info("extractAuthTables-sql:{}",sql);
        return extractAuthTables(sql, dbType, false,dbName,viewAuthTables);
    }

    public void writeResult(String fileName, String result) {
        try {
            File file = new File(fileName);
            if (!file.exists()) {
                file.createNewFile();
            }
            FileUtils.writeStringToFile(new File(fileName), result, Charset.defaultCharset(), false);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

//    public static void main(String[] args) {
//        String sql = args[0];
//        String dbType = args[1];
//        String auth_param = args[2];
//        String fileName = args[3];
//        SqlAuthentication sqlAuthentication = new SqlAuthentication();
//        boolean result = sqlAuthentication.authenticate(sql, dbType, auth_param);
//        sqlAuthentication.writeResult(fileName, String.valueOf(result));
//    }

}
